from langchain.agents import AgentType, initialize_agent, AgentExecutor, create_tool_calling_agent
from tool.neo4j_tool import Neo4jTool
from tool.document_tool import DocumentTool
from tool.search_tool import SearchTool
from tool.show_tool import ShowTool
from model.my_chat_model import ChatModel
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder


# 创建一个聊天的智能体
def create_agent(question):
    # 1 获取大模型
    chat = ChatModel()
    llm = chat.get_line_model()
    # 2 创建工具
    tools = [Neo4jTool(), DocumentTool(), SearchTool(), ShowTool()]
    # 3 创建提示模板
    prompt = ChatPromptTemplate.from_messages([
        ("system", """
         你是一个智能的AI查询助手，你可以使用以下几个工具
         1 get_neo4j_tool 主要用于查询知识图谱的数据，例如人物，公司的关系数据
         2 get_document_tool 主要用于查询文档类型数据，例如公司详细说明
         3 get_search_tool 主要用于当get_neo4j_tool和get_document_tool查询不到数据时才使用
         4 get_show_tool 当用于生成图表或者可视化图时使用
         如果用户的提问的问题是有关于人物，公司的一些基本数据，请执行工具 get_neo4j_tool,
         如果是有关于公司财务或者详细类型的数据，请执行工具 get_document_tool,
         如果get_neo4j_tool和get_document_tool查询不到数据时,执行get_search_tool,
         如果需要生成某公司图表或可视化，就执行工具get_show_tool,
         如果没有需要图表可视化等，就不执行get_show_tool工具
         
         图表回复格式（强制约束，绝对禁止图片链接）
        1. 回复必须包含 3 部分：工具日志 + 文字说明 + ECharts JSON 配置；
        2. ECharts JSON 必须用 ```json 和 ``` 严格包裹（前端将通过此标识提取配置）；
        3. 绝对禁止出现任何图片链接（如 ![](xxx.png)），禁止用图片占位；
        4. 若未按此格式输出，将判定为无效回复，需重新生成。
        5. 绝对禁止输出 ```json 和 ``` 这两个标记
        """),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),  # 工具调用
    ])
    # 创建工具调用agentx
    agent = create_tool_calling_agent(llm, tools, prompt)
    # 创建智能体执行器
    agent_executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,

    )
    rs = agent_executor.invoke({"input": question})
    print(rs)
    return rs["output"]


if __name__ == '__main__':
    create_agent("给出茅台的十大股东信息")
